package com.nbsaas.boot.sql;

import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Function;

import java.util.ArrayList;
import java.util.List;

public class SelectFieldNameExtractor2 {
    public static List<String> getFieldNames(String sql) {
        List<String> fields = new ArrayList<>();
        try {
            Statement stmt = CCJSqlParserUtil.parse(sql);
            if (stmt instanceof Select) {
                PlainSelect select = (PlainSelect) ((Select) stmt).getSelectBody();

                for (SelectItem item : select.getSelectItems()) {
                    Expression expr = item.getExpression(); // 从 SelectItem 获取 Expression

                    if (expr instanceof Column) {
                        fields.add(((Column) expr).getColumnName());
                    }
                    else if (expr instanceof Function) {
                        Function func = (Function) expr;
                        if (func.getParameters() != null) {
                            func.getParameters().getExpressions().forEach(param -> {
                                if (param instanceof Column) {
                                    fields.add(((Column) param).getColumnName());
                                }
                            });
                        }
                    }
                    // 其他表达式类型可在此扩展处理
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return fields;
    }

    public static void main(String[] args) {
        String sql1 = "SELECT id, name, age FROM users WHERE age > 20";
        String sql2 = "SELECT u.id, u.name AS userName, COUNT(o.id) total FROM users u " +
                      "JOIN orders o ON u.id = o.user_id GROUP BY u.id, u.name";

        System.out.println(getFieldNames(sql1)); // 输出: [id, name, age]
        System.out.println(getFieldNames(sql2)); // 输出: [id, name, id]
    }
}
